# Image denoising using MRF model
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
from skimage import util


def main():
    # Read in image
    im = Image.open('./images/lena512.bmp')
    im = np.array(im)
    im = (im - im.min()) / (im.max() - im.min())
    img = np.where(im > 0.5, 1, 0)  # convert to binary image

    # Add noise
    # noisy = im.copy()
    # noise = np.random.rand(M, N)
    # ind = where(noise < 0.2)
    # noisy[ind] = 1 - noisy[ind]

    im_noisy = util.random_noise(im, mode='s&p', seed=0)
    # mode: gaussian, localvar, poisson, salt, pepper, s&p, speckle
    img_noise = np.where(im_noisy > 0.5, 1, 0)  # convert to binary image

    out = MRF_denoise(img_noise)

    plt.figure()
    # gray()
    plt.subplot(131)
    plt.imshow(img, cmap=plt.gray())
    plt.axis('off')
    plt.title('Original Image')

    plt.subplot(132)
    plt.imshow(img_noise, cmap=plt.gray())
    plt.axis('off')
    plt.title('Noisy Image')

    plt.subplot(133)
    plt.imshow(out, cmap=plt.gray())
    plt.axis('off')
    plt.title('Denoised Image')
    plt.show()


def MRF_denoise(noisy):
    # Start MRF
    (M, N) = noisy.shape
    y_old = noisy
    y = np.zeros((M, N))

    while SNR(y_old, y) > 0.01:
        print('SNR: {}'.format(SNR(y_old, y)))
        for i in range(M):
            for j in range(N):
                index = neighbor(i, j, M, N)

                a = cost(1, noisy[i, j], y_old, index)
                b = cost(0, noisy[i, j], y_old, index)

                if a > b:
                    y[i, j] = 1
                else:
                    y[i, j] = 0
        y_old = y
    print('End >>> SNR: {}'.format(SNR(y_old, y)))
    return y


def SNR(A, B):
    if A.shape == B.shape:
        return np.sum(np.abs(A - B)) / A.size
    else:
        raise Exception("Two matrices must have the same size!")


def delta(a, b):
    if a == b:
        return 1
    else:
        return 0


def neighbor(i, j, M, N):
    # find correct neighbors
    if i == 0 and j == 0:
        neighbors = [(0, 1), (1, 0)]
    elif i == 0 and j == N - 1:
        neighbors = [(0, N - 2), (1, N - 1)]
    elif i == M - 1 and j == 0:
        neighbors = [(M - 1, 1), (M - 2, 0)]
    elif i == M - 1 and j == N - 1:
        neighbors = [(M - 1, N - 2), (M - 2, N - 1)]
    elif i == 0:
        neighbors = [(0, j - 1), (0, j + 1), (1, j)]
    elif i == M - 1:
        neighbors = [(M - 1, j - 1), (M - 1, j + 1), (M - 2, j)]
    elif j == 0:
        neighbors = [(i - 1, 0), (i + 1, 0), (i, 1)]
    elif j == N - 1:
        neighbors = [(i - 1, N - 1), (i + 1, N - 1), (i, N - 2)]
    else:
        neighbors = [(i - 1, j), (i + 1, j), (i, j - 1), (i, j + 1), (i - 1, j - 1), (i - 1, j + 1), (i + 1, j - 1),
                     (i + 1, j + 1)]
    return neighbors


def cost(y, x, y_old, index):
    alpha = 1
    beta = 10
    return alpha * delta(y, x) + beta * sum([delta(y, y_old[i]) for i in index])


if __name__ == "__main__":
    main()
